/* PRT_Exporter Plug-in

 Copyright (c) 2010, Steven Caron
 All rights reserved.

 Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met:

 1.Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer.
 2.Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or
   other materials provided with the distribution.
 3.Neither the name of the copyright holder nor the names of its contributors may be used to endorse or promote products derived from this software
   without specific prior written permission.

 THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
 LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
 HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
 LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON
 ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE
 USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE
 
 This plugin implements Krakatoa's particle cache format, outlined here...
 http://software.primefocusworld.com/software/support/krakatoa/prt_file_format.php
 
 All questions, comments, bug reports should be sent to tools at steven-caron.com
 
 Special thanks to Helge Mathee, Adam Ferrall, and Mark Wiebe for helping me write this plugin!    
*/

#include <stdio.h>
#include <assert.h>
#include <string.h>

#include <xsi_application.h>
#include <xsi_context.h>
#include <xsi_pluginregistrar.h>
#include <xsi_status.h>
#include <xsi_primitive.h>
#include <xsi_geometry.h>
#include <xsi_iceattribute.h>
#include <xsi_iceattributedataarray.h>

#include <zlib.h>

using namespace XSI; 
using namespace MATH;

const ULONG CHUNKSIZE = 10000;

XSIPLUGINCALLBACK CStatus XSILoadPlugin( PluginRegistrar& in_reg ){
	in_reg.PutAuthor( L"Steven Caron" );
	in_reg.PutName( L"PRT Exporter Plug-in" );
	in_reg.PutEmail( L"tools@steven-caron.com" );
	in_reg.PutURL( L"http://www.steven-caron.com" );
	in_reg.PutVersion( 1, 0 );
	
	in_reg.RegisterConverterEvent( L"siOnPRTExport", siOnCustomFileExport, L"prt" );

	return CStatus::OK;
}

XSIPLUGINCALLBACK CStatus XSIUnloadPlugin( const PluginRegistrar& in_reg ){
	CString strPluginName;
	strPluginName = in_reg.GetName();
	Application().LogMessage( strPluginName + L" has been unloaded.", siVerboseMsg );
	return CStatus::OK;
}

// a helper function to determine if a ICE attibute is defined
bool AttributeIsDefined( CRefArray in_attributes, CString in_name ){
	for(long i=0; i < in_attributes.GetCount(); i++ ){
		if( ICEAttribute(in_attributes[i]).GetName().IsEqualNoCase( in_name ) )
			return true;
	}
	
	return false;
}

// Callback for the siOnPRTExport event.
XSIPLUGINCALLBACK CStatus siOnPRTExport_OnEvent( CRef& in_ctxt ){
	
	//get the context object
	Context ctxt( in_ctxt );
	
	// force the frame to be LONG, we can't allow sub frames here.
	LONG frameNb = CValue( ctxt.GetAttribute( L"Frame" ) );

	// retrieve the geometry from the context
	Geometry geo = Primitive( CRef( ctxt.GetAttribute( L"Target") ) ).GetGeometry( (double)frameNb );

	CString fileName = CString( ctxt.GetAttribute( L"FileName" ) );
	for( LONG i = fileName.Length() - 6 ; i >= 0; i-- ){
		
		if( fileName[i] == '.' || fileName[i] == '_' ){
			
			// rebuild the filename based on the delimiter positions
			CString newFileName;
			for( LONG j=0; j <= i; j++ ){
				newFileName += fileName[j];
			}
			newFileName += CValue( frameNb ).GetAsText() + L".prt";
			fileName = newFileName;
			break;
		}
	}

	// open the file
	FILE * file = fopen( fileName.GetAsciiString(), "wb" );
	
	if ( NULL != file ){
		
		//write the file header
		//magic number for the PRT file format
		char buf[] = { 192, 'P', 'R', 'T', '\r', '\n', 26, '\n' };
		fwrite( buf, sizeof(char), sizeof(buf) / sizeof(char), file ); 
		
		//length of header file
		int intTemp = 56;
		fwrite( &intTemp, 4, 1, file );
		
		//(32 bytes) A human readable signature null-terminated string describing the file, currently "Extensible Particle Format".
		char epf[32] = "Extensible Particle Format";
		fwrite( epf, sizeof( epf ), 1, file );
		
		//(4 bytes)  A 32 bit int indicating version (Has value 1).
		intTemp = 1;
		fwrite( &intTemp, 4, 1, file );
		
		//get the point position and point velocity attributes
		ICEAttribute posAttr = geo.GetICEAttributeFromName( L"PointPosition" );
		ICEAttribute velAttr = geo.GetICEAttributeFromName( L"PointVelocity" );
		ICEAttribute colorAttr = geo.GetICEAttributeFromName( L"Color" );

		//(8 bytes)  A 64 bit int indicating particle count.
		ULONG nbParticles = posAttr.GetElementCount();
		__int64 countTemp = nbParticles;
		fwrite( &countTemp, 8, 1, file );
		
		//reserved section bytes
		intTemp = 4;
		fwrite( &intTemp, 4, 1, file );
		
		//number of attributes
		intTemp = 4;
		fwrite( &intTemp, 4, 1, file );
		
		//standard attribute definition structure (44)
		intTemp = 44;
		fwrite( &intTemp, 4, 1, file );
		
		CRefArray allAttributes = geo.GetICEAttributes();
		
		//write point position and velocity
		if( AttributeIsDefined( allAttributes, L"PointPosition" ) && AttributeIsDefined( allAttributes, L"PointVelocity" ) && AttributeIsDefined( allAttributes, L"Color" )){
			
			//particle data header
			//position attribute
			//(32 bytes) A null-terminated string indicating the attribute name
			char pos[32] = "Position";
			fwrite( pos, 32, 1, file );
			
			//attribute type... 4 for float32
			intTemp = 4;
			fwrite( &intTemp, 4, 1, file );

			//number of data values each particle has for this attribute, ie 3 for vectors and colors, 1 for scalars and integers
			intTemp = 3;
			fwrite( &intTemp, 4, 1 ,file );			
			
			//attribute offset... 0 for the first vector attribute
			//this offset should be computed properly based on the nb of attributes and its data type
			intTemp = 0;
			fwrite( &intTemp, 4, 1, file );
			
			//velocity attribute
			char vel[32] = "Velocity";
			fwrite( vel, 32, 1, file );
			
			intTemp = 4;
			fwrite( &intTemp, 4, 1, file );

			intTemp = 3;
			fwrite( &intTemp, 4, 1 ,file );			
			
			//attribute offset... 12 for the second vector attribute
			intTemp = 12;
			fwrite( &intTemp, 4, 1, file );
			
			//color attribute
			char color[32] = "Color";
			fwrite( color, 32, 1, file );
			
			intTemp = 4;
			fwrite( &intTemp, 4, 1, file );

			intTemp = 3;
			fwrite( &intTemp, 4, 1 ,file );			
			
			//attribute offset... 24 for the third vector attribute
			intTemp = 24;
			fwrite( &intTemp, 4, 1, file );

			//density attribute
			char density[32] = "Density";
			fwrite( density, 32, 1, file );

			intTemp = 4;
			fwrite( &intTemp, 4, 1, file );
			
			//scalar/single float32 value
			intTemp = 1;
			fwrite( &intTemp, 4, 1, file );

			//attribute offset... 36 (three previous vectors)
			intTemp = 36;
			fwrite( &intTemp, 4, 1, file );

			// also test to see if the attributes contain any data...
			CICEAttributeDataArray<CVector3f> posTestData;
			CICEAttributeDataArray<CVector3f> velTestData;
			CICEAttributeDataArray<CColor4f> colorTestData;
			posAttr.GetDataArrayChunk( 0, 1, posTestData );
			velAttr.GetDataArrayChunk( 0, 1, velTestData );
			colorAttr.GetDataArrayChunk( 0, 1, colorTestData );
			
			//setup vector objects and rotation for conversion to 3dsmax coord space
			CVector3f posVector;
			CVector3f velVector;
			CMatrix3f rotMatrix( 1.0, 0.0, 0.0, 0.0, 0.0, -1.0, 0.0, 1.0, 0.0 ); //90 degree rotation
			rotMatrix.TransposeInPlace();

			if( posTestData.GetCount() > 0 && velTestData.GetCount() > 0 && colorTestData.GetCount() > 0 ){
				
				//hardcoding PARTICLESIZE for now
				const int PARTICLESIZE = 40;
				int level = -1; //-1 compression level is zlib default
				int ret, flush;
				unsigned have;
				z_stream strm;
				unsigned char in[PARTICLESIZE];
				unsigned char out[PARTICLESIZE];
				
				//allocate deflate state
				strm.zalloc = Z_NULL;
				strm.zfree = Z_NULL;
				strm.opaque = Z_NULL;
				ret = deflateInit(&strm, level);
				if (ret != Z_OK)
					return ret;

				// read the data in chunks of 1000 elements
				for( ULONG i=0; i < nbParticles; i += CHUNKSIZE ){
					
					CICEAttributeDataArray<CVector3f> posData;
					CICEAttributeDataArray<CVector3f> velData;
					CICEAttributeDataArray<CColor4f> colorData;
					
					ULONG upperLimit = nbParticles - i < CHUNKSIZE ? nbParticles - i : CHUNKSIZE;

					posAttr.GetDataArrayChunk( i, upperLimit, posData );
					velAttr.GetDataArrayChunk( i, upperLimit, velData );
					colorAttr.GetDataArrayChunk( i, upperLimit, colorData );
					for( ULONG j=0; j < upperLimit; j++ ){
						
						//rotate the vectors
						posVector.Set( posData[j].GetX(), posData[j].GetY(), posData[j].GetZ() );
						posVector.MulByMatrix3InPlace( rotMatrix );
						velVector.Set( velData[j].GetX(), velData[j].GetY(), velData[j].GetZ() );
						velVector.MulByMatrix3InPlace( rotMatrix );

						float floatTemp[10];
						//float floatAlpha;
						posVector.Get( floatTemp[0], floatTemp[1], floatTemp[2] ); //position
						velVector.Get( floatTemp[3], floatTemp[4], floatTemp[5] ); //velocity
						colorData[j].GetAsRGBA( floatTemp[6], floatTemp[7], floatTemp[8], floatTemp[9] ); //color and density in alpha
						
						memcpy((void *)in, (const void *)floatTemp, PARTICLESIZE);
                        strm.avail_in = PARTICLESIZE;
						
						//to figure out what the state of deflate() should be this iteration
						flush = i + j == nbParticles - 1 ? Z_FINISH : Z_NO_FLUSH;
						strm.next_in = in;
						
						//operate until avail_out is actually empty, it seems deflate never gets all of avail_out
						//even in our case where we are working with a really small amount of data each deflate() request
						do{
							strm.avail_out = PARTICLESIZE; //avail_out must be at least 0.1% larger than avail_in plus 12 bytes, according to zlib docs
							strm.next_out = out;
							ret = deflate(&strm, flush);
							assert(ret != Z_STREAM_ERROR);
							have = PARTICLESIZE - strm.avail_out;
							if(fwrite(out, 1, have, file) != have || ferror(file)){
								(void)deflateEnd(&strm);
								return Z_ERRNO;
							}
						}while(strm.avail_out == 0); //make sure avail_out is actually empty
						assert(strm.avail_in == 0);
					
					}	
				}
				assert(ret == Z_STREAM_END); //make sure zlib returned Z_STREAM_END
				
				// clean up and return
				(void)deflateEnd(&strm);
			}
			else{
				Application().LogMessage( CRef( ctxt.GetAttribute( L"Target") ).GetAsText() + L" doesn't contain all the proper attributes on frame " + CValue(frameNb).GetAsText() + L"!", siErrorMsg );
				fclose(file);	
			}
		}	
		// close the file
		fclose(file);
	}

	return CStatus::OK;
}